package tk.hongbo.stream.coder

import androidx.camera.core.ImageProxy


/**
 * Created by HongboZhao on 2022/1/28.
 */
object YUV {

    /**
     * 把NV21转换成NV12
     * 不转换则出现编码后颜色翻转，蓝色变为红色，红色变为蓝色
     */
    private fun nv21toNV12(nv21: ByteArray): ByteArray {
        val size = nv21.size
        val nv12 = ByteArray(size)
        val len = size * 4 / 6
        System.arraycopy(nv21, 0, nv12, 0, len)
        var i = len
        while (i < size - 1) {
            nv12[i] = nv21[i + 1]
            nv12[i + 1] = nv21[i]
            i += 2
        }
        return nv12
    }

    /**
     * 从YUV_420_888获取NV12
     */
    fun getMediaCodecData(image: ImageProxy): ByteArray {
        return rotateYUV420Degree90(nv21toNV12(yuv420ToNv21(image)), image.width, image.height, 90)
    }

    /**
     * YUV_420_888转NV21
     *
     * @param image CameraX ImageProxy
     * @return byte array
     */
    private fun yuv420ToNv21(image: ImageProxy): ByteArray {
        val planes = image.planes
        val yBuffer = planes[0].buffer
        val uBuffer = planes[1].buffer
        val vBuffer = planes[2].buffer
        val ySize = yBuffer.remaining()
        val uSize = uBuffer.remaining()
        val vSize = vBuffer.remaining()
        val size = image.width * image.height
        val nv21 = ByteArray(size * 3 / 2)
        yBuffer[nv21, 0, ySize]
        vBuffer[nv21, ySize, vSize]
        val u = ByteArray(uSize)
        uBuffer[u]

        //每隔开一位替换V，达到VU交替
        var pos = ySize + 1
        for (i in 0 until uSize) {
            if (i % 2 == 0) {
                nv21[pos] = u[i]
                pos += 2
            }
        }
        image.close()
        return nv21
    }

    fun rotateYUV420Degree90(input: ByteArray, width: Int, height: Int, rotation: Int): ByteArray {
        val frameSize = width * height
        val qFrameSize = frameSize / 4
        val output = ByteArray(frameSize + 2 * qFrameSize)
        val swap = rotation == 90 || rotation == 270
        val yflip = rotation == 90 || rotation == 180
        val xflip = rotation == 270 || rotation == 180
        for (x in 0 until width) {
            for (y in 0 until height) {
                var xo = x
                var yo = y
                var w = width
                var h = height
                var xi = xo
                var yi = yo
                if (swap) {
                    xi = w * yo / h
                    yi = h * xo / w
                }
                if (yflip) {
                    yi = h - yi - 1
                }
                if (xflip) {
                    xi = w - xi - 1
                }
                output[w * yo + xo] = input[w * yi + xi]
                val fs = w * h
                val qs = fs shr 2
                xi = xi shr 1
                yi = yi shr 1
                xo = xo shr 1
                yo = yo shr 1
                w = w shr 1
                h = h shr 1
                val ui = fs + (w * yi + xi) * 2
                val uo = fs + (w * yo + xo) * 2
                val vi = ui + 1
                val vo = uo + 1
                output[uo] = input[ui]
                output[vo] = input[vi]
            }
        }
        return output
    }

    private fun rotateYUV420Degree90(
        data: ByteArray,
        imageWidth: Int,
        imageHeight: Int
    ): ByteArray? {
        val yuv = ByteArray(480 * 640 * 3 / 2)
        // Rotate the Y luma
        var i = 0
        for (x in 0 until imageWidth) {
            for (y in imageHeight - 1 downTo 0) {
                yuv[i] = data[y * imageWidth + x]
                i++
            }
        }
        // Rotate the U and V color components
        i = imageWidth * imageHeight * 3 / 2 - 1
        var x = imageWidth - 1
        while (x > 0) {
            for (y in 0 until imageHeight / 2) {
                yuv[i] = data[imageWidth * imageHeight + y * imageWidth + x]
                i--
                yuv[i] = data[imageWidth * imageHeight + y * imageWidth + (x - 1)]
                i--
            }
            x = x - 2
        }
        return yuv
    }

    /**
     * 此处为顺时针旋转旋转90度
     *
     * @param data        旋转前的数据
     * @param imageWidth  旋转前数据的宽
     * @param imageHeight 旋转前数据的高
     * @return 旋转后的数据
     */
    private fun rotateYUV420Degree901(
        data: ByteArray,
        imageWidth: Int,
        imageHeight: Int
    ): ByteArray {
        val yuv = ByteArray(imageWidth * imageHeight * 3 / 2)
        // Rotate the Y luma
        var i = 0
        for (x in 0 until imageWidth) {
            for (y in imageHeight - 1 downTo 0) {
                yuv[i] = data[y * imageWidth + x]
                i++
            }
        }
        // Rotate the U and V color components
        i = imageWidth * imageHeight * 3 / 2 - 1
        var x = imageWidth - 1
        while (x > 0) {
            for (y in 0 until imageHeight / 2) {
                yuv[i] = data[imageWidth * imageHeight + y * imageWidth + x]
                i--
                yuv[i] = data[imageWidth * imageHeight + y * imageWidth + (x - 1)]
                i--
            }
            x = x - 2
        }
        return yuv
    }

    private fun rotateYUV420Degree180(
        data: ByteArray,
        imageWidth: Int,
        imageHeight: Int
    ): ByteArray {
        val yuv = ByteArray(imageWidth * imageHeight * 3 / 2)
        var i = 0
        var count = 0
        i = imageWidth * imageHeight - 1
        while (i >= 0) {
            yuv[count] = data[i]
            count++
            i--
        }
        i = imageWidth * imageHeight * 3 / 2 - 1
        i = imageWidth * imageHeight * 3 / 2 - 1
        while (i >= imageWidth
            * imageHeight
        ) {
            yuv[count++] = data[i - 1]
            yuv[count++] = data[i]
            i -= 2
        }
        return yuv
    }

    /**
     * 此处为顺时针旋转270
     *
     * @param data        旋转前的数据
     * @param imageWidth  旋转前数据的宽
     * @param imageHeight 旋转前数据的高
     * @return 旋转后的数据
     */
    private fun rotateYUV420Degree270(
        data: ByteArray,
        imageWidth: Int,
        imageHeight: Int
    ): ByteArray {
        val yuv = ByteArray(imageWidth * imageHeight * 3 / 2)

        // Rotate the Y luma
        var i = 0
        for (x in imageWidth - 1 downTo 0) {
            for (y in 0 until imageHeight) {
                yuv[i] = data[y * imageWidth + x]
                i++
            }
        } // Rotate the U and V color components
        i = imageWidth * imageHeight
        var x = imageWidth - 1
        while (x > 0) {
            for (y in 0 until imageHeight / 2) {
                yuv[i] = data[imageWidth * imageHeight + y * imageWidth + (x - 1)]
                i++
                yuv[i] = data[imageWidth * imageHeight + y * imageWidth + x]
                i++
            }
            x = x - 2
        }
        return yuv
    }

    //镜像
    private fun Mirror(yuv_temp: ByteArray, w: Int, h: Int) {
        var i: Int
        var j: Int
        var a: Int
        var b: Int
        var temp: Byte
        //mirror y
        i = 0
        while (i < h) {
            a = i * w
            b = (i + 1) * w - 1
            while (a < b) {
                temp = yuv_temp[a]
                yuv_temp[a] = yuv_temp[b]
                yuv_temp[b] = temp
                a++
                b--
            }
            i++
        }
        //mirror u
        var uindex = w * h
        i = 0
        while (i < h / 2) {
            a = i * w / 2
            b = (i + 1) * w / 2 - 1
            while (a < b) {
                temp = yuv_temp[a + uindex]
                yuv_temp[a + uindex] = yuv_temp[b + uindex]
                yuv_temp[b + uindex] = temp
                a++
                b--
            }
            i++
        }
        //mirror v
        uindex = w * h / 4 * 5
        i = 0
        while (i < h / 2) {
            a = i * w / 2
            b = (i + 1) * w / 2 - 1
            while (a < b) {
                temp = yuv_temp[a + uindex]
                yuv_temp[a + uindex] = yuv_temp[b + uindex]
                yuv_temp[b + uindex] = temp
                a++
                b--
            }
            i++
        }
    }

}